# -*- coding: utf-8 -*-
import sys   
import os
import warnings
import numpy as np
import pandas as pd

# 如果你有自定义 pymagic 路径，加入 sys.path
dirs = ['/database/home/duansizhang/pywave']
for d in dirs:
    if d not in sys.path:
        sys.path.insert(0, d)

from pymagic.tools.ecgdllinterface import DataSource, EcgAnalyzer
from hrvanalysis import (
    remove_outliers, interpolate_nan_values, remove_ectopic_beats,
    get_time_domain_features, get_frequency_domain_features,
    get_geometrical_features, get_sampen
)
import EntropyHub

warnings.filterwarnings("ignore")

# ========== 参数配置 ==========
raw_file = '/database/private/mgcdb/raw/0a0b34b8220644848209d7a199afd6fc_1699056000000.raw'
out_dir  = '/database/home/duansizhang/hrv_predict/result'
fs       = 250       # 采样率 (Hz)
win_sec  = 300       # 窗口长度：5分钟

os.makedirs(out_dir, exist_ok=True)

def extract_features(nn_seg):
    """对给定 NN 段计算时域、频域、几何和非线性特征（nbeats, ApEn 除外）"""
    if len(nn_seg) < 3:
        return None
    feats = {}
    td = get_time_domain_features(nn_seg)
    feats.update({k: float(td[k]) if td[k] is not None else np.nan for k in td})
    fd = get_frequency_domain_features(nn_seg)
    feats.update({k: float(fd[k]) if fd[k] is not None else np.nan for k in fd})
    geo = get_geometrical_features(nn_seg)
    feats.update({k: float(geo[k]) if geo[k] is not None else np.nan for k in geo})
    feats['nbeats'] = len(nn_seg)
    feats['ApEn']   = np.nan  # 后面再补
    return feats

# 1. 载入并检测 R 波
hdata   = DataSource(raw_file, 3)
res     = EcgAnalyzer(hdata)
rpos_all= np.array(res['rpos'], dtype=int)
anntype = np.array(res['anntype'], dtype=int)
rpos    = rpos_all[anntype == 1]

if len(rpos) < 2:
    raise RuntimeError("R 波太少，无法计算 RR")

# 2. RR 预处理
times     = rpos / fs
rr_ms     = np.diff(rpos) * 1000.0 / fs
rr_times  = times[1:]

rr_clean  = remove_outliers(rr_intervals=rr_ms, low_rri=300, high_rri=2000)
rr_interp = interpolate_nan_values(rr_intervals=rr_clean, interpolation_method='linear')
nn_beats  = remove_ectopic_beats(rr_intervals=rr_interp, method='malik')
nn_interp = interpolate_nan_values(rr_intervals=nn_beats)
mask      = ~np.isnan(nn_interp)
nn_array  = np.array(nn_interp)[mask]
nn_times  = rr_times[mask]

if len(nn_array) < 3:
    raise RuntimeError("有效 NN 段太短")

# 3. 全局 + 窗口特征提取
rows = []

# 全局
gf = extract_features(nn_array)
rows.append({
    'label': 'global',
    'window_start': 0.0,
    'window_end': float(nn_times[-1]),
    'timestamp': 0.0,
    **gf
})

# 窗口（从 0 到末尾，每个窗口 win_sec，**不做数量上限**）
starts = np.arange(0, float(nn_times[-1]), win_sec)
for ws in starts:
    we  = ws + win_sec
    seg = nn_array[(nn_times >= ws) & (nn_times < we)]
    ft  = extract_features(seg)
    if ft is None:
        continue
    # 仅当段内beat数足够多时，计算非线性特征
    if len(seg) > 10:
        sm = get_sampen(seg)
        ft.update({k: float(sm[k]) for k in sm})
        try:
            ap, _ = EntropyHub.ApEn(seg)
            ft['ApEn'] = float(ap[2])
        except:
            ft['ApEn'] = np.nan

    rows.append({
        'label': 'window',
        'window_start': ws,
        'window_end': we,
        'timestamp': ws + win_sec/2,
        **ft
    })

df_out = pd.DataFrame(rows)
# 按需调整列顺序
cols = ['label', 'window_start', 'window_end', 'timestamp'] + [c for c in df_out.columns if c not in ['label','window_start','window_end','timestamp']]
df_out = df_out[cols]

# 4. 保存
basename = os.path.basename(raw_file).replace('.raw','_rrfeature.csv')
out_path = os.path.join(out_dir, basename)
df_out.to_csv(out_path, index=False, encoding='utf-8-sig')

print(f"已生成特征文件：{out_path}")